import numpy as np
import zarr
import argparse
import os
import pickle
from pathlib import Path
from vis_pcd import PointArrow, catpcd
from scalePcd import scalePcdD
from save_file import save_info
from scipy.spatial.transform import Rotation
import open3d as o3d
def transform_arrows(arrows_array, center_point, offset_rpy):
    arrows_array=np.array(arrows_array)
    center_point=np.array(center_point)
    offset_rpy=np.array(offset_rpy)
    assert arrows_array.shape[1] == 6, f"arrows_array shape error: {arrows_array.shape}"
    assert center_point.shape == (3,), f"center_point shape error: {center_point.shape}"
    assert offset_rpy.shape == (3,), f"offset_rpy shape error: {offset_rpy.shape}"
    roll, pitch, yaw = offset_rpy
    center_x, center_y, center_z = center_point
    rotation_matrix = np.array([
        [np.cos(yaw) * np.cos(pitch), np.cos(yaw) * np.sin(pitch) * np.sin(roll) - np.sin(yaw) * np.cos(roll), np.cos(yaw) * np.sin(pitch) * np.cos(roll) + np.sin(yaw) * np.sin(roll)],
        [np.sin(yaw) * np.cos(pitch), np.sin(yaw) * np.sin(pitch) * np.sin(roll) + np.cos(yaw) * np.cos(roll), np.sin(yaw) * np.sin(pitch) * np.cos(roll) - np.cos(yaw) * np.sin(roll)],
        [-np.sin(pitch), np.cos(pitch) * np.sin(roll), np.cos(pitch) * np.cos(roll)]
    ])
    transformed_arrows = np.zeros_like(arrows_array)
    for arrow_index in range(arrows_array.shape[0]):
        arrow_xyz = arrows_array[arrow_index, :3] - center_point
        arrow_rpy = arrows_array[arrow_index, 3:]
        transformed_xyz = np.dot(rotation_matrix, arrow_xyz) + center_point
        transformed_rpy = arrow_rpy + offset_rpy
        transformed_arrows[arrow_index, :3] = transformed_xyz
        transformed_arrows[arrow_index, 3:] = transformed_rpy
    return transformed_arrows
def state6d_to_arrows(states, start_color=[0, 0, 1], end_color=[1, 0, 0]):
    pointcloud = np.zeros((0, 6))
    num_timesteps = len(states)
    for timestep_idx, timestep_data in enumerate(states):
        xyz_rpy = timestep_data[:6]
        time_ratio = timestep_idx / (num_timesteps - 1) if num_timesteps > 1 else 0
        color = [start_color[color_idx] * (1 - time_ratio) + end_color[color_idx] * time_ratio for color_idx in range(3)]
        arrow_points = PointArrow(xyz_rpy, color, 50)
        pointcloud = catpcd(pointcloud, arrow_points)
    return pointcloud
def test_transform(base_arrow_xyzrpy, offset_rpy, original_data_states):
    assert base_arrow_xyzrpy.shape == (6,), f"base_arrow_xyzrpy shape error: {base_arrow_xyzrpy.shape}"
    assert offset_rpy.shape == (3,), f"offset_rpy shape error: {offset_rpy.shape}"
    base_arrow_pcd = PointArrow(base_arrow_xyzrpy, [1, 0, 0], 50)
    original_arrows_pcd = state6d_to_arrows(original_data_states)
    combined_pcd = catpcd(base_arrow_pcd, original_arrows_pcd)
    center_point = base_arrow_xyzrpy[:3]
    transformed_arrows = transform_arrows(original_data_states, center_point, offset_rpy)
    transformed_arrows_pcd = state6d_to_arrows(transformed_arrows)
    final_pcd = catpcd(combined_pcd, transformed_arrows_pcd)
    save_info("pcd", "test_transform_result", final_pcd)
    return transformed_arrows
if __name__ == "__main__":
    base_arrow = np.array([0, 0, 0, 0, 0, 0])
    offset_rpy = np.array([30, 30, 120])
    original_states = np.linspace([0.1, 0.1, 0.1, 0, 0, 0], [0.3, 0.1, 0.1, 0, 0, 0], 5)
    result = test_transform(base_arrow, offset_rpy, original_states)